import os
import math as m
import numexpr as ne
from numpy import arctan2, sqrt
import open3d as o3d
import numpy as np
import copy
import sys
import cv2
import yaml


def draw_registration_result(source, target, transformation):
    source_temp = copy.deepcopy(source)
    target_temp = copy.deepcopy(target)
    source_temp.paint_uniform_color([1, 0.706, 0])
    target_temp.paint_uniform_color([0, 0.651, 0.929])
    source_temp.transform(transformation)
    o3d.io.write_point_cloud("source_temp.ply", source_temp)
    o3d.visualization.draw_geometries([source_temp, target_temp])


def tfa():
    SOURCE_PCD = "/home/lei/code/ply/0.ply"
    TARGET_PCD = "/home/lei/桌面/viz_elas_setero/build/all0.ply"
    # SOURCE_PCD = sys.argv[1]
    # TARGET_PCD = sys.argv[2]
    source = o3d.io.read_point_cloud(SOURCE_PCD)
    target = o3d.io.read_point_cloud(TARGET_PCD)

    # trans_init = np.asarray([[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0],
    #                         [0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 1.0]])

    draw_registration_result(source, target, trans_init)

# fs2 = cv2.FileStorage('/home/u20/桌面/Aruco-Marker/build/a.yaml', cv2.FileStorage_READ)
# # p0 = fs2.getNode('p0').mat()
# # p1 = fs2.getNode('p1').mat()
# # s0 = o3d.io.read_point_cloud("cp0.ply")
# # s1 = o3d.io.read_point_cloud("cp1.ply")
# # s1.transform(p1)
# # s0.transform(p0)
# axis_pcd = o3d.geometry.TriangleMesh.create_coordinate_frame(size=0.5, origin=[0, 0, 0])
# pcds=[axis_pcd]
# for i in range(1):
#     transformation = fs2.getNode('p'+str(i)).mat()
#     # transformation=np.linalg.inv(transformation)
#     print(transformation)
#     source_temp = o3d.io.read_point_cloud(f"cp{i}.ply")
#     # source_temp.transform(transformation)
#     pcds.append(source_temp)
# o3d.visualization.draw_geometries(pcds)


dpps = float(sys.argv[2])


def filter(point_stack):
    pts = np.array(point_stack.points)
    z = pts[:, 2]
    x = pts[:, 0]
    # field1 = z < 20
    ss = np.where(z >= dpps)[0]
    print(ss)
    # point_stack.colors=  point_stack.colors[field2]
    inlier_cloud = pcd.select_by_index(list(ss))
    pts = np.array(inlier_cloud.points)[:, 0]
    ss = np.where(pts > -0.5)[0]
    inlier_cloud = inlier_cloud.select_by_index(list(ss))

    pts = np.array(inlier_cloud.points)[:, 1]
    ss = np.where(pts > -0.5)[0]
    inlier_cloud = inlier_cloud.select_by_index(list(ss))

    # print(ss)
#  and pts[:, 0]<1 and pts[:, 0]>-1

    # outlier_cloud = pcd.select_by_index(list(ss), invert=True)
    # outlier_cloud.paint_uniform_color([0.0,0.0,  1.0])
    return inlier_cloud, 0


filepath = sys.argv[1]
axis_pcd = o3d.geometry.TriangleMesh.create_coordinate_frame(size=1.0, origin=[
                                                             0, 0, 0])
# -*-coding:utf-8 -*-
test_data_dir = './'
point_cloud_file_name = filepath+'/mesh.ply'
point_cloud_file_path = os.path.join(test_data_dir, point_cloud_file_name)
# 读取点云
pcd = o3d.io.read_point_cloud(point_cloud_file_path)
pcdo, _ = filter(pcd)


def cal_aabb(pcd):
    print("->正在计算点云轴向最小包围盒...")
    aabb = pcd.get_axis_aligned_bounding_box()
    # print(aabb.get_min_bound())
    # print(aabb.get_max_bound())
    sasd = aabb.get_max_bound()-aabb.get_min_bound()
    print(np.array(aabb.get_box_points()))
    aabb.color = (1, 0, 0)
    print("->正在计算点云最小包围盒...")
    obb = pcd.get_oriented_bounding_box()
    print(np.array(obb.get_box_points()))
    obb.color = (0, 1, 0)
    # o3d.visualization.draw_geometries(
    # [pcd, aabb, obb], window_name="wechat 394467238")
    o3d.io.write_point_cloud(filepath+"/plant.pcd", pcdo, write_ascii=True)
    os.system(f"./py/area {filepath} 0.002 ")

    py_object = dict({'length': float(sasd[0]), 'width': float(
        sasd[1]), 'height': float(sasd[2])})
    file = open(f"{filepath}/lhw.yml", 'w')
    yaml.dump(py_object, file)
    file.close()

    return sasd


sasd = cal_aabb(pcdo)

# o3d.visualization.draw_geometries([axis_pcd, pcdo])


# def seg(pcd):
#     # 平面分割
#     plane_model, inliers = pcd.segment_plane(distance_threshold=0.05,
#     ransac_n=3,
#     num_iterations=1000)
#     # 模型参数
#     [a, b, c, d] = plane_model
#     print(f"Plane equation: {a:.2f}x + {b:.2f}y + {c:.2f}z + {d:.2f} = 0")


#     # 可视化
# op,ip=seg(pcd)
# o3d.visualization.draw_geometries([op,ip ])
